# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
R* 算法工具函数模块

本模块提供了 R* 深度思考算法所需的各种工具函数，主要包括：
- 数学答案提取和格式化处理
- 提示词构建和包装
- 步骤结果解析和观察包装
- 数学等价性判断和答案验证
- LaTeX 格式处理和文本清理

核心功能:
1. 答案提取: 从 LaTeX 格式文本中提取数学答案
2. 提示词处理: 构建符合 R* 格式的提示词模板
3. 结果解析: 解析模型输出为结构化的步骤结果
4. 等价判断: 使用多种评估工具判断数学答案的等价性
5. 格式清理: 处理各种数学表达式的格式问题

主要组件:
- 文本处理函数: remove_text_box, extract_boxed_answer, extract_math_answer
- 提示词函数: rstar_prompt_wrap, rstar_obs_wrap, rstar_step_result_unwrap
- 等价判断函数: rstar_equiv, math_equiv, is_multi_choice
- 辅助工具函数: remove_single_dollar, any_condition
"""
from __future__ import annotations  # 启用延迟注解，支持前向引用类型提示
from typing import List, Dict, Any, Optional, Type, Tuple, Union  # 类型提示工具
from math_evaluation import is_equiv  # 数学等价性判断工具
from rstar_deepthink.prompts.prompt_rstar import PROMPT_RSTAR  # R* 提示词模板
from rstar_deepthink.tools.python_tool import PythonInterpreter  # Python 解释器工具
from rstar_deepthink.constants import *  # 项目常量定义
from utils.math_equal import math_equal  # 数学等式判断工具
from utils.checker import check_one_answer  # 答案检查工具
from utils.util import equiv, strip_string, choice_answer_clean  # 通用工具函数


def remove_text_box(text: Optional[str]) -> Optional[str]:
    """
    移除 LaTeX 文本中的 \text{} 标记
    
    该函数用于处理 LaTeX 格式的数学表达式，移除其中的 \text{} 包装，
    保留文本内容或在特殊情况下完全移除。主要用于清理数学答案中的
    格式标记，使答案更易于比较和验证。
    
    参数:
        text: 包含 LaTeX \text{} 标记的文本字符串
        
    返回:
        Optional[str]: 移除 \text{} 标记后的文本，如果输入为 None 则返回 None
        
    处理逻辑:
    1. 查找最后一个 \text{ 标记的位置
    2. 使用栈匹配对应的右括号 }
    3. 提取 \text{} 内部的文本内容
    4. 如果内容是 "and"，则完全移除该标记
    5. 否则保留内部文本，移除 \text{} 包装
    
    示例:
        >>> remove_text_box("x = 5 \\text{meters}")
        "x = 5 meters"
        >>> remove_text_box("a \\text{and} b")
        "a  b"
    """
    if text is None:
        return None
    start = text.find(r"\text{")
    if start == -1:
        return text
    end = None
    stack = []
    answer = text[start:]
    for i, c in enumerate(answer):
        if c == "{":
            stack.append(i)
        elif c == "}":
            start_text = stack.pop()
            if len(stack) == 0:
                end_text = i
                break
    in_text_string = text[start + start_text + 1 : start + end_text]

    if in_text_string.strip() == "and":
        ex_text = text[:start] + text[start + end_text + 1 :]
    else:
        ex_text = (
            text[:start]
            + text[start + start_text + 1 : start + end_text].strip()
            + text[start + end_text + 1 :]
        )
    return ex_text.strip()


def extract_boxed_answer(text: Optional[str], debug: bool = False) -> Optional[str]:
    """
    从 LaTeX 文本中提取 \boxed{} 包装的答案
    
    该函数专门用于从数学推理文本中提取最终答案，这些答案通常用
    LaTeX 的 \boxed{} 命令包装。函数会查找最后一个 \boxed{} 标记，
    并提取其中的内容作为最终答案。
    
    参数:
        text: 包含 LaTeX \boxed{} 标记的文本字符串
        debug: 是否启用调试模式，当括号不匹配时输出错误信息
        
    返回:
        Optional[str]: 提取的答案内容，如果没有找到 \boxed{} 则返回原文本，
                      如果输入为 None 则返回 None
                      
    处理逻辑:
    1. 查找最后一个 \boxed{ 标记的位置（使用 rfind 确保获取最后一个）
    2. 使用栈算法匹配对应的右括号 }
    3. 提取括号内的完整内容
    4. 处理嵌套括号的情况
    5. 在调试模式下报告括号不匹配的错误
    
    示例:
        >>> extract_boxed_answer("The answer is \\boxed{42}")
        "42"
        >>> extract_boxed_answer("\\boxed{x^2 + y^2 = z^2}")
        "x^2 + y^2 = z^2"
    """
    if text is None:
        return None
    start = text.rfind(r"\boxed{")
    if start == -1:
        return text
    end = None
    stack = []
    answer = text[start:]
    for i, c in enumerate(answer):
        if c == "{":
            stack.append(i)
        elif c == "}":
            start = stack.pop()  # \boxed start{
            if len(stack) == 0:
                end = i  # \boxed end}
                break
    if end is None and debug:
        print("brack not closing", answer)
        return None
    return answer[start + 1 : end]

# 无效答案的标识符，用于标记提取失败的情况
INVALID_ANS = "[invalid]"

def extract_math_answer(answer: str) -> str:
    """
    从数学推理文本中提取并清理最终答案
    
    该函数是答案提取的主要入口点，结合了多种文本处理技术来
    从复杂的数学推理文本中提取干净的最终答案。它处理各种格式
    问题，包括句号、换行符、LaTeX 标记等。
    
    参数:
        answer: 包含数学答案的原始文本字符串
        
    返回:
        str: 清理后的数学答案，如果提取失败则返回 INVALID_ANS
        
    处理流程:
    1. 按 '.\n' 分割文本，取第一部分（移除后续解释）
    2. 移除末尾的句号（如果存在）
    3. 提取 \boxed{} 包装的内容
    4. 移除 \text{} 标记
    5. 清理多余的空白字符
    6. 异常处理：任何错误都返回 INVALID_ANS
    
    异常处理:
        任何在处理过程中发生的异常都会被捕获，函数返回 INVALID_ANS
        
    示例:
        >>> extract_math_answer("The answer is \\boxed{42}.\n")
        "42"
        >>> extract_math_answer("\\boxed{\\text{meters}} = 5")
        "meters = 5"
    """
    try:
        ans = answer
        extract_ans_temp = ans.split('.\n')[0]
        extract_ans_temp = extract_ans_temp.strip()
        if len(extract_ans_temp) > 0 and extract_ans_temp[-1] == '.':
            extract_ans = extract_ans_temp[0:-1]
        else:
            extract_ans = extract_ans_temp
        extract_ans = extract_ans.strip()
        extract_ans = remove_text_box(extract_boxed_answer(extract_ans))
    except:
        extract_ans = INVALID_ANS
    return extract_ans


# Python 解释器工具的字符串描述，用于提示词中的工具说明
python_tool_string = f"{PythonInterpreter().name}: {PythonInterpreter().description}"
# Python 解释器工具的名称，用于动作识别
python_tool_name = PythonInterpreter().name
    

def rstar_prompt_wrap(
    question: str, 
    partial_solution: str,
    config,
) -> str:
    """
    构建 R* 格式的提示词
    
    该函数是 R* 算法的核心提示词构建器，负责将问题、部分解答和配置
    组合成符合 R* 格式要求的完整提示词。它集成了示例、格式说明和
    当前的求解进度，为语言模型提供结构化的输入。
    
    参数:
        question: 要解决的数学问题
        partial_solution: 当前的部分解答（可能为空）
        config: 配置对象，包含步骤分隔符等设置
        
    返回:
        str: 构建完成的 R* 格式提示词
        
    构建流程:
    1. 获取配置中的步骤分隔符
    2. 创建 R* 提示词模板实例
    3. 随机选择演示示例
    4. 根据示例数量生成相应的前缀说明
    5. 组合格式说明、示例前缀和示例内容
    6. 添加问题输入和后缀模板
    7. 如果有部分解答，则追加到提示词末尾
    
    提示词结构:
    - 格式说明（format_instructions）
    - 示例前缀（example_prefix）
    - 演示示例（rstar_examples）
    - 问题输入（inputs）
    - 部分解答（partial_solution）
    
    注意:
        函数会根据示例数量动态调整前缀说明的单复数形式
    """
    # 获取配置中的步骤分隔符
    step_delim = config.step_delim
    # 初始化R-STAR提示词模板
    prompt_pot = PROMPT_RSTAR(config)
    # 构建输入字符串，包含问题内容和步骤分隔符
    inputs = f"{question}{step_delim}"  

    # 从提示词模板中随机获取示例
    rstar_examples = prompt_pot.random_examples()
    
    # 根据示例数量设置示例前缀文本
    if len(rstar_examples) > 1:
        example_prefix = "The following are %d demonstration examples." % len(rstar_examples)
    elif len(rstar_examples) == 1:
        example_prefix = "The following is a demonstration example."

    # 获取提示词格式说明
    format_instructions = prompt_pot.pot_format_instructions
    
    # 如果有示例，构建包含格式说明、示例前缀和示例的提示词
    if len(rstar_examples) > 0:
        prompt = step_delim.join([format_instructions, example_prefix, *rstar_examples, ""])
    else:
        # 没有示例时只包含格式说明
        prompt = step_delim.join([format_instructions, ""])
    
    # 如果提示词为空，使用提示词模板的后缀格式
    if prompt.strip() == "":
        prompt = step_delim.join([prompt_pot.pot_suffix.format(input=inputs)])
    else:
        # 否则将后缀格式添加到现有提示词后
        prompt = step_delim.join([prompt, prompt_pot.pot_suffix.format(input=inputs)])
    
    # 如果有部分解决方案，将其添加到提示词中
    if partial_solution:
        prompt = "".join([prompt, partial_solution])
    
    # 返回最终构建的提示词
    return prompt + ""


def rstar_obs_wrap(observation: str) -> str:
    """
    包装代码执行的观察结果
    
    该函数将代码执行的输出结果用标准的输出标记包装，
    使其符合 R* 格式的要求。这样可以在后续处理中
    清晰地识别和解析代码执行的结果。
    
    参数:
        observation: 代码执行后的观察结果字符串
        
    返回:
        str: 用 OUTPUT 和 OUTPUT_END 标记包装的观察结果
        
    格式:
        {OUTPUT}观察结果内容{OUTPUT_END}
        
    用途:
        - 标准化代码执行结果的格式
        - 便于后续的文本解析和处理
        - 与 R* 提示词格式保持一致
    """
    return f"{OUTPUT}{observation}{OUTPUT_END}"


def rstar_step_result_unwrap(
    text: str,
) -> Tuple[str, Dict[str, str]]:
    """
    解析 R* 格式的步骤结果
    
    该函数是 R* 算法的核心解析器，负责将语言模型的输出文本
    解析为结构化的步骤结果。它能够识别最终答案和中间步骤，
    并将它们转换为标准的动作-输入格式。
    
    参数:
        text: 语言模型输出的原始文本
        
    返回:
        Tuple[str, Dict[str, str]]: 包含原始文本和解析结果的元组
            - 第一个元素：原始文本
            - 第二个元素：解析结果字典，包含以下键：
                - "action": 动作类型
                - "action_input": 动作输入内容
                - "final_answer": 最终答案（如果有）
                
    解析逻辑:
    1. 检查文本中是否包含 "boxed" 关键词
    2. 如果包含，则认为是最终答案：
       - 提取数学答案并设置为 final_answer
       - action 和 action_input 保持为空
    3. 如果不包含，则认为是中间步骤：
       - 设置 action 为 "python_interpreter"
       - 将整个文本设置为 action_input
       
    注意:
        该函数假设所有非最终答案的文本都是 Python 代码
    """
    parser_result = {
        "action": "",
        "action_input": "",
        "final_answer": "",
    }
    #if ANSWER_END in text or "boxed" in text:
    if "boxed" in text:
        parser_result["final_answer"] = extract_math_answer(text)
        return text, parser_result
    else:
        parser_result["action"] = "python_interpreter"
        parser_result["action_input"] = text
        return text, parser_result


def is_multi_choice(answer: str) -> bool:
    """
    判断答案是否为多选题格式
    
    检查给定的答案字符串是否只包含多选题的选项字母（A-E）。
    这个函数用于识别多选题答案，以便进行相应的格式处理。
    
    参数:
        answer: 要检查的答案字符串
        
    返回:
        bool: 如果答案只包含 A-E 字母则返回 True，否则返回 False
        
    示例:
        >>> is_multi_choice("ABC")
        True
        >>> is_multi_choice("A1B")
        False
        >>> is_multi_choice("")
        True
    """
    for c in answer:
        if c not in ["A", "B", "C", "D", "E"]:
            return False
    return True


def remove_single_dollar(s: Union[str, List[str], None]) -> Union[str, None]:
    """
    移除字符串两端的单个美元符号
    
    该函数用于清理 LaTeX 数学表达式中的美元符号包装。
    如果输入是列表，则取第一个元素进行处理。
    
    参数:
        s: 要处理的字符串、字符串列表或 None
        
    返回:
        Union[str, None]: 移除美元符号后的字符串，如果输入为空则返回原值
        
    处理逻辑:
    1. 如果输入为空，直接返回
    2. 如果输入是列表，取第一个元素
    3. 检查是否同时以 $ 开头和结尾
    4. 如果是，则移除两端的美元符号
    
    示例:
        >>> remove_single_dollar("$x^2$")
        "x^2"
        >>> remove_single_dollar(["$abc$", "def"])
        "abc"
        >>> remove_single_dollar("$incomplete")
        "$incomplete"
    """
    if not s:
        return s
    if isinstance(s, list):
        s = s[0]
    if s.startswith("$") and s.endswith("$"):
        s = s[1:-1]
    return s


def any_condition(conditions: List[bool]) -> bool:
    """
    检查条件列表中是否有任何一个为真
    
    这是 Python 内置 any() 函数的简单包装，用于提高代码的可读性。
    
    参数:
        conditions: 布尔值条件列表
        
    返回:
        bool: 如果任何一个条件为 True 则返回 True，否则返回 False
        
    示例:
        >>> any_condition([False, True, False])
        True
        >>> any_condition([False, False, False])
        False
    """
    return any(conditions)


def rstar_equiv(gt: str, pred: str) -> bool:
    """
    R* 算法的核心等价性判断函数
    
    该函数集成了多个开源评估工具，每个工具都有自己的判断逻辑和
    在处理特殊情况（如 LaTeX、单位等）方面的优势。通过多重验证
    策略来确保数学答案等价性判断的准确性和鲁棒性。
    
    参数:
        gt: 标准答案（ground truth）
        pred: 预测答案
        
    返回:
        bool: 如果两个答案等价则返回 True，否则返回 False
        
    验证策略:
    1. 直接字符串比较
    2. 使用多个评估工具进行双向比较（gt vs pred 和 pred vs gt）
       - math_equal: 数学等式判断
       - is_equiv: 等价性检查
       - check_one_answer: 答案验证
    3. 字符串清理后再次比较（处理格式问题）
    4. 多选题特殊处理：
       - 单选答案 vs 非选择格式：清理预测答案
       - 多选答案 vs 非选择格式：提取选项字母
    
    特殊处理:
    - 对于 college-math 和 omni-math 数据集，会交换 gt 和 pred 的位置
      进行比较，因为发现这些数据集中部分问题的标准答案质量较低
    - 多选题答案会进行格式标准化处理
    
    异常处理:
        任何比较过程中的异常都会被捕获并打印，函数返回 False
        
    注意:
        该函数是 R* 算法答案验证的核心，直接影响最终的准确率评估
    """
    # 集成了多个开源评估工具，每个工具都有自己的判断逻辑和
    # 在处理特殊情况（如 LaTeX、单位等）方面的优势
    gt = str(gt)
    pred = str(pred)
    try:
        if gt == pred:
            return True
        
        # 对于 college-math 和 omni-math，需要交换 pred 和 gt 的位置
        # 因为我们发现在 college-math 等基准测试的一小部分问题中，标准答案的质量相对较低
        if any(
            func(x, y) for func in [math_equal, is_equiv, check_one_answer] for x, y in [(gt, pred), (pred, gt)]
        ):
            return True
        # 针对 college-math 等的特殊处理
        gt_strip, pred_strip = strip_string(gt), strip_string(pred)
        if any(
            func(x, y) for func in [math_equal, is_equiv, check_one_answer] for x, y in [(gt_strip, pred_strip), (pred_strip, gt_strip)]
        ):
            return True

        # 针对选择题的处理
        if gt in ["A", "B", "C", "D", "E"] and pred not in ["A","B","C","D","E"]:
            pred = choice_answer_clean(pred)
            if math_equal(gt, pred):
                return True
        elif is_multi_choice(gt) and not is_multi_choice(pred):
            pred = "".join(
                    [c for c in pred if c in ["A", "B", "C", "D", "E"]]
                )
            if math_equal(gt, pred):
                return True
    except Exception as e:
        print("maroi_equiv error")
        print(e)
        pass
    return False
        

def rstar_step_result_unwrap_wo_code(
    text: str,
) -> Tuple[str, Dict[str, str]]:
    """
    无代码模式的 R* 格式步骤结果解析
    
    该函数是专门为无代码 MCTS 模式设计的解析器，负责将语言模型的输出文本
    解析为结构化的步骤结果。与原有的解析函数不同，它不处理 Python 代码执行，
    只关注纯思考推理和最终答案。
    
    参数:
        text: 语言模型输出的原始文本
        
    返回:
        Tuple[str, Dict[str, str]]: 包含原始文本和解析结果的元组
            - 第一个元素：原始文本
            - 第二个元素：解析结果字典，包含以下键：
                - "action": 动作类型（保持为空，无代码执行）
                - "action_input": 动作输入内容（保持为空）
                - "final_answer": 最终答案（如果有）
                
    解析逻辑:
    1. 检查文本中是否包含 "boxed" 关键词（最终答案标记）
    2. 如果包含，则认为是最终答案：
       - 提取数学答案并设置为 final_answer
       - action 和 action_input 保持为空
    3. 如果不包含，则认为是中间思考步骤：
       - 所有字段都保持为空，表示普通思考文本
    4. 保留现有的动作解析逻辑以确保多模态兼容性
    
    特点:
    - 不设置 action 为 "python_interpreter"，避免代码执行
    - 完全复用现有的最终答案提取逻辑
    - 保持与原有 action 机制的兼容性
    """
    parser_result = {
        "action": "",
        "action_input": "",
        "final_answer": "",
    }
    
    # 检查是否包含最终答案标记
    if "boxed" in text:
        parser_result["final_answer"] = extract_math_answer(text)
        return text, parser_result
    else:
        # 对于无代码模式，普通文本不设置任何 action
        # 这样会被识别为普通思考文本，不会触发代码执行
        return text, parser_result


def math_equiv(grt: Union[str, List[str]], prd: str) -> bool:
    """
    数学等价性判断的高级接口函数
    
    该函数是 rstar_equiv 的高级包装器，支持处理多个标准答案的情况。
    当标准答案是列表时，只要预测答案与其中任何一个等价即认为正确。
    这在处理有多个正确答案的数学问题时特别有用。
    
    参数:
        grt: 标准答案，可以是单个字符串或字符串列表
        prd: 预测答案字符串
        
    返回:
        bool: 如果预测答案与任何一个标准答案等价则返回 True，否则返回 False
        
    处理逻辑:
    1. 如果标准答案是列表：
       - 遍历列表中的每个标准答案
       - 使用 rstar_equiv 逐一比较
       - 只要有一个匹配就返回 True
       - 如果都不匹配则返回 False
    2. 如果标准答案是单个字符串：
       - 直接调用 rstar_equiv 进行比较
    
    用途:
    - 处理有多个正确表达形式的数学问题
    - 提供统一的答案验证接口
    - 简化上层调用代码的复杂性
    
    示例:
        >>> math_equiv(["1/2", "0.5"], "0.5")
        True
        >>> math_equiv("42", "42")
        True
        >>> math_equiv(["A", "B"], "C")
        False
    """
    prd = (prd)
    if isinstance(grt, list):
        for g in grt:
            if rstar_equiv(g, prd):
                return True
        return False
    else:
        return rstar_equiv(grt, prd)
